include("geomatry.jl")
include("operators.jl")
using DifferentialEquations, Plots, BenchmarkTools, Sundials
using ModelingToolkit

function rhs(u,p,t)
    D, Φ, h, df, Np, K, α = p   
  
    df[1,1] = (α/2.0 - 1)*(u[1,1]-u[end,end])
    df[1,2:K] = (α/2.0 - 1)*(u[1,2:K]-u[end,1:K-1])
    df[end,end] = α/2.0*(u[end,end]-u[1,1])
    df[end,1:K-1]=α/2.0*(u[end,1:K-1]-u[1,2:K])
    return -2.0/h*(D*u .- Φ*df)
end

function rhs1(du,u,p,t)
    D, Φ, h, df, Np, K, α = p   

    df[1,1] = (α/2.0 - 1)*(u[1,1]-u[end,end])
    df[1,2:K] = (α/2.0 - 1)*(u[1,2:K]-u[end,1:K-1])
    df[end,end] = α/2.0*(u[end,end]-u[1,1])
    df[end,1:K-1]=α/2.0*(u[end,1:K-1]-u[1,2:K])
    
    du .= -2.0/h*(D*u - Φ*df)  
end

function explicit(p, k)
    """
    Solve the discrete ODE by explicit time-martching methods.
    """
    P = p         # order of polynomials
    K = k         # number of elements
    Np = P+1      # number of points at each element 

    # calculate the element-wise operators
    r = LegendreGaussLobatto(P)   # solution points at [-1,1]
    V = Vandermonde(P, r)         # Vandermonde matrix at [-1,1]
    D = DiffMatrix(P, r, V)       # Differential matrix D = M^-1*S
    Φ = V*V'                      # correct matrix Φ = M^-1
    
    # get the geomatry informations
    h, x = Geomatry1D(0, 2π, P, K, r) # global solution points x and h
    
    # initialize the solution
    u0 = sin.(x)

    # definations for DifferentialEquations package
    df = zeros(Np, K)
    tspan = [0, 1.0]
    p = (D, Φ, h, df, Np, K, 0)

    # define the ode problem
    # prob = ODEProblem(rhs, u0, tspan, p)
    prob = ODEProblem(rhs1, u0, tspan, p)

    uf = solve(prob,Tsit5(),save_everystep=false)
    
    # calculate the error
    ue = sin.(x .- tspan[2])
    ϵ = norm(ue - uf[2])
    print("the error is:", ϵ)
end

function implicit(p, k)
    """
    Solve the discrete ODE by implicit time-martching methods.
    """
    P = p         # order of polynomials
    K = k         # number of elements
    Np = P+1      # number of points at each element 

    # calculate the element-wise operators
    r = LegendreGaussLobatto(P)   # solution points at [-1,1]
    V = Vandermonde(P, r)         # Vandermonde matrix at [-1,1]
    D = DiffMatrix(P, r, V)       # Differential matrix D = M^-1*S
    Φ = V*V'                      # correct matrix Φ = M^-1
    
    # get the geomatry informations
    h, x = Geomatry1D(0, 2π, P, K, r) # global solution points x and h
    
    # initialize the solution
    u0 = sin.(x)

    # definations for DifferentialEquations package
    df = zeros(Np, K)
    tspan = [0, 1.0]
    
    # define the ODE problems
    p = (D, Φ, h, df, Np, K, 0)
    # prob = ODEProblem(rhs, u0, tspan, p)
    prob = ODEProblem(rhs1, u0, tspan, p)

    # solve the ODEs
    @time sol1 = solve(prob,CVODE_BDF(linear_solver=:GMRES),save_everystep=false)
    @time sol2 = solve(prob,TRBDF2(autodiff=false),progress=true,save_everystep=false,save_start=false)
    # calculate the error
    ue = sin.(x .- tspan[2])
    ϵ1 = norm(ue - sol1[2])
    print("the error is:", ϵ1)
end